"""
与登录相关的控制器
@Time: 2021/1/22 18:20
@Author:
@File: login_controller.py
"""
from fastapi import Depends, HTTPException, status, APIRouter
from fastapi.security import OAuth2PasswordRequestForm
from sqlalchemy.orm import Session
from jose import jwt
from passlib.context import CryptContext
from pydantic import BaseModel, Field
from typing import Optional, Union
from datetime import datetime, timedelta
from redis import StrictRedis
import redis

from config import user_token_conf
from models.dto import RESTfulModel
from dependencies import get_mysql_db, get_redis, verify_token, __oauth2_scheme
from models.po import mysql_po
from mapper.mysql_mapper import user_mapper


login_router = APIRouter(tags=['login'])


class Token(BaseModel):
    code: int = 0
    access_token: str = Field(description='本次登录的token')
    token_type: str = Field(default='Bearer', description='token的类型，统一为 Bearer')
    student_name: str = Field(description='学生姓名')


pwd_context = CryptContext(schemes=['bcrypt'], deprecated='auto')


def get_pwd_hash(pwd):
    return pwd_context.hash(pwd)


def authenticate_user(mysql_db: Session, username: str, password: str) -> Union[bool, mysql_po.User]:
    """
    验证用户合法性
    :return: 若验证成功则返回 User 对象；若验证失败则返回 False
    """
    user: mysql_po.User = user_mapper.get_user(mysql_db, username)
    if not user:
        return False
    if not pwd_context.verify(password, user.hashed_password):
        return False
    return user


def create_access_token(data: dict, expires_delta: Optional[timedelta] = None):
    to_encode = data.copy()
    if expires_delta:
        expire = datetime.utcnow() + expires_delta
    else:
        expire = datetime.utcnow() + timedelta(minutes=15)
    to_encode.update({"exp": expire})
    encoded_jwt = jwt.encode(to_encode, user_token_conf.SECRET_KEY, algorithm=user_token_conf.ALGORITHM)
    return encoded_jwt


@login_router.post("/token",
                   response_model=Token,
                   summary='登录接口，获取 token',
                   description="""采用OAuth2.0认证协议，登录时获取用户的token，
                                使用 x-www-form-urlencoded 格式, 
                                在 Request Body 提交 username 和 password 即可获得本次用户登录的token,
                                并会被缓存到 Redis 中保持一定的时间段""")
async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends(),
                                 mysql_db: Session = Depends(get_mysql_db),
                                 redis_db: StrictRedis = Depends(get_redis)):
    user = authenticate_user(mysql_db, form_data.username, form_data.password)  # 查询数据库进行用户验证
    if not user:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Incorrect username or password",
            headers={"WWW-Authenticate": "Bearer"},
        )
    # 根据 username 生成 token
    access_token_expires = timedelta(minutes=user_token_conf.ACCESS_TOKEN_EXPIRE_MINUTES)
    access_token = create_access_token(
        data={"sub": user.username}, expires_delta=access_token_expires
    )

    redis_db.set(access_token, user.username, ex=user_token_conf.ACCESS_TOKEN_EXPIRE_MINUTES * 60)  # 设置 redis_db 缓存

    return Token(
        code=0,
        access_token=access_token,
        token_type='Bearer',
        student_name=user.student_name
    )


@login_router.get("/test",
                  response_model=RESTfulModel,
                  dependencies=[Depends(verify_token), ],
                  summary='token 验证测试',
                  tags=['test'])
async def test():
    return RESTfulModel(code=0, data='success')



@login_router.get("/logout",
                  response_model=RESTfulModel[str],
                  dependencies=[Depends(verify_token), ],
                  summary='用户登出',
                  description="清除掉存储于 Redis 的该用户 token，使 token 失效")
async def logout(token: str = Depends(__oauth2_scheme),
                 redis: StrictRedis = Depends(get_redis)):
    try:
        redis.delete(token)
    except Exception as e:
        return RESTfulModel(code=-1, data=f'Redis 出现异常: {e}')
    return RESTfulModel(code=0, data='success')
